# USE CASE: 
# Takes in a dataset with ID (unique_id), weight (idvdwt_crsec_coreimgrt), the independent variable (neverNewOldStmp2) and all the depending variables 

# Note: 
# 1) Generate the bootstrap weights using generate_bs_weights.py
# 2) Not each bootstrap sample will converge. You can exclude that sample and generate a new set of weights and rerun the analysis.

# Import necessary libraries
library(readstata13)
library(Matrix)
require(glmnet)
library(MASS)


# NOTE: Set these parameters to run desired analysis
num_bootstrap_samples <- 50
input_dta_file_path <- "Put the path of the dta data file of interest."
bs_weight_path_prefix <- "Put the prefix of the path to each set of bootstrap weights." # i.e. "E:/bs_weights_"
output_coef_prefix <- "Put the prefix to what the coefficients are going to saved as." # i.e. "E:/lasso_coefs_"
output_mse_mins <- "Put the path to file where the mse mins for each boostrap run is saved."
output_lambda_mins <- "Put the path to file where the lambdas corresponding to the mse mins for each boostrap run i saved."

# Read in data file
d <- read.dta13(input_dta_file_path)

# Drop any observations with missing values
d1 <- d[complete.cases(d),]

# Create vectors that keep track of the lambdas and minimum mean squared errors for each 
# bootstrap sample
lambdas <- c()
mse_mins <- c()

# For every bootstrap sample
for (i in c(1:num_bootstrap_samples)) {

  # Read in the bootstrap weights
  sweights <- read.csv(paste(bs_weight_path_prefix, i, ".csv", sep=""), header=FALSE)
  sweights <- sweights$V1
  ids <- d1$unique_id
  y <- d1$neverNewOldStmp2

  # Keep only dependent variables and transform into a matrix
  x <- subset(d1, select=-c(unique_id, idvdwt_crsec_coreimgrt, neverNewOldStmp2))
  x1 <- data.matrix(x)
  x1 <- Matrix(x1, sparse = TRUE)
  
  # Fitting the model using the botstrap weights
  cv.lasso <- cv.glmnet(x1,y, family = 'multinomial', weights = sweights, parallel = TRUE)

  # Extract coefficients corresponding to the minimum mean squared error
  coefMins <- coef(cv.lasso, s=cv.lasso$lambda.min)
  coef1 <- matrix(coefMins[[1]])
  coef2 <- matrix(coefMins[[2]])
  coef3 <- matrix(coefMins[[3]])

  # Extract the minimum mean squared error
  mse.min <- cv.lasso$cvm[cv.lasso$lambda == cv.lasso$lambda.min]

  # Append the minimum mean squared error and corresponding lambda
  mse_mins <- c(mse_mins, mse.min)
  lambdas <- c(lambdas, cv.lasso$lambda.min)

  # Write all variables to a file for prediction purposes
  write.table(cbind(coef1, coef2, coef3), paste0(output_coef_prefix,i,".txt"))
}

# Write the MSE mins and corresponding lambdas to files
write.csv(mse_mins, output_mse_mins)
write.csv(lambdas, output_lambda_mins)
